# -*- coding: utf-8 -*-

"""
计算几个常用的图谱变更集：完整变更集、变更报告、Neo4j增量发版变更集、ES增量发版变更集
"""

from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division

from typing import Set, AnyStr, Dict, List, Tuple

from operator import add

from kgpipeline.job import KgJob, JobConfigError
from kgpipeline.jsonutil import entity_json_decode, entity_json_encode
from kgpipeline.cnschema import KgCnSchema
from kgpipeline.calc.diff import DEFAULT_IGNORE_KEYS_LINK, DEFAULT_IGNORE_KEYS, key_diff, diff_handler_neo4j,\
    diff_handler_es
from kgpipeline.calc.rdd_diff import rdd_diff_by_id, handle_diff

from pyspark import RDD, Broadcast
from pyspark.sql import DataFrame
from kgpipeline.sparkutil import to_all_string_dataframe, to_text_dataframe

import logging

logger = logging.getLogger(__name__)


# diff_type, id, main_type, property, v1, v2
# ADD, _____, MutualFund, fundManager, N/A, 20
# DEL, _____, MutualFund, fundManager, 20, N/A
# UPD_ADD, _____, MutualFund, fundManager, N/A, 20
# UPD_DEL, _____, MutualFund, fundManager, 20, N/A
# UPD_UPD, _____, MutualFund, fundManager, 20, 21
def diff_handler_detail(dict1, dict2, ignore, schema):  # type: (Dict, Dict, Set[AnyStr], KgCnSchema) -> List[Tuple]
    res = []
    if dict1 is None:
        entity_id = dict2["@id"]
        main_type = "|".join(schema.get_main_types(dict2["@type"]))
        for key, value in dict2.items():
            res.append(("DEL", entity_id, main_type, key, value, None))
    elif dict2 is None:
        entity_id = dict1["@id"]
        main_type = "|".join(schema.get_main_types(dict1["@type"]))
        for key, value in dict1.items():
            res.append(("ADD", entity_id, main_type, key, None, value))
    else:
        key_remove, key_add, key_update = key_diff(dict1, dict2, ignore)
        main_type = "|".join(schema.get_main_types(dict1["@type"]))
        entity_id = dict1["@id"]
        for key in key_remove:
            res.append(("UPD_DEL", entity_id, main_type, key, dict2[key], None))
        for key in key_add:
            res.append(("UPD_ADD", entity_id, main_type, key, None, dict1[key]))
        for key in key_update:
            res.append(("UPD_UPD", entity_id, main_type, key, dict2[key], dict1[key]))
    return res


# diff report of types
# diff_type, main_type, main_type_zh, property, property_zh, count
# ADD, MutualFund, N/A, 100
# DEL, MutualFund, N/A, 20
# UPD_ADD, MutualFund, fundManager, 100
# UPD_DEL,
# UPD_UPD,
def generate_diff_report(diff, broadcast_schema):  # type: (RDD, Broadcast) -> DataFrame
    fields = ["diff_type", "main_type", "main_type_zh", "property", "property_zh", "count"]
    report = diff.map(lambda x: ((x[0], x[2], x[3]), 1)).reduceByKey(add).map(
        lambda x: (x[0][0],
                   x[0][1], broadcast_schema.value.get_type_name_zh(x[0][1]),
                   x[0][2], broadcast_schema.value.get_type_name_zh(x[0][2]),
                   x[1])
    )
    return to_all_string_dataframe(report, fields).orderBy("diff_type", "main_type", "property")


class KgDiffJob(KgJob):
    """
    calculate several diffs

    config includes:
    - schema: schema pkl path
    - entity_ignore: list of entities to ignore. optional, default is DEFAULT_IGNORE_KEYS
    - link_ignore: list of links to ignore. optional, default is DEFAULT_IGNORE_KEYS_LINK
    - disable_report: whether to disable diff report. optional, default is False
    - disable_detail: whether to disable detail report. optional, default is False
    - disable_neo4j: whether to disable neo4j report. optional, default is False
    - disable_es: whether to disable es report. optional, default is False
    """
    type_key = "diff"

    def __init__(self, config, reader):
        super(KgDiffJob, self).__init__(config, reader)
        self._schema = self.config["schema"]

    def process(self, inputs):
        """
        kg diff process
        """
        logger.info("Start KgDiffJob: {}...".format(self.name))

        if len(inputs) != 4:
            raise JobConfigError(
                "Expect 4 inputs: current entity.json and link.json"
                " vs other entity.json and link.json while we have {0}! Abort!".format(
                    len(inputs)))

        entity_ignore = set(self.config.get("entity_ignore", []))
        link_ignore = set(self.config.get("link_ignore", []))
        disable_report = self.config.get("disable_report", False)
        disable_detail = self.config.get("disable_detail", False)
        disable_neo4j = self.config.get("disable_neo4j", False)
        disable_es = self.config.get("disable_es", False)

        entity_ignore.update(DEFAULT_IGNORE_KEYS)
        link_ignore.update(DEFAULT_IGNORE_KEYS_LINK)

        rdd_list = [df.rdd.map(lambda x: entity_json_decode(x["value"])) if df is not None else None for df, _ in
                    inputs]
        current_entity, current_link, other_entity, other_link = rdd_list

        if current_entity is None or other_entity is None:
            raise JobConfigError("Current and other entity input must be specified!")

        entity_diff = rdd_diff_by_id(current_entity, other_entity)

        link_diff = None
        if current_link is not None or other_link is not None:
            if current_link is None or other_link is None:
                raise JobConfigError("Either current or other link input is not specified!")
            link_diff = rdd_diff_by_id(current_link, other_link)

        diff_full = handle_diff(entity_diff, entity_ignore, diff_handler_detail, self._schema)
        if link_diff is not None:
            diff_full = diff_full.union(handle_diff(link_diff, link_ignore, diff_handler_detail, self._schema))

        res = []
        if not disable_detail:
            diff = to_all_string_dataframe(diff_full, ["diff_type", "id", "main_type", "property", "v1", "v2"])
            res.append(diff)
        else:
            res.append(None)

        if not disable_report:
            diff_report = generate_diff_report(diff_full, self._schema)
            res.append(diff_report)
        else:
            res.append(None)

        if not disable_neo4j:
            entity_fix = handle_diff(entity_diff, entity_ignore, diff_handler_neo4j)
            res.append(to_text_dataframe(entity_fix, entity_json_encode))
            if link_diff is not None:
                link_fix = handle_diff(link_diff, link_ignore, diff_handler_neo4j)
                res.append(to_text_dataframe(link_fix, entity_json_encode))
        else:
            res.append(None)
            res.append(None)

        if not disable_es:
            es_diff = handle_diff(entity_diff, entity_ignore, diff_handler_es)
            res.append(to_text_dataframe(es_diff, entity_json_encode))
        else:
            res.append(None)

        return res


KgJob.register("diff", KgDiffJob)
